/-
Copyright (c) 2025 Madison Crim, Aaron Liu, Justus Springer, Junyan Xu. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Madison Crim, Aaron Liu, Justus Springer, Junyan Xu
-/
module

public import Mathlib.Algebra.Module.PID
public import Mathlib.Algebra.MvPolynomial.Funext
public import Mathlib.Algebra.Polynomial.Module.AEval
public import Mathlib.FieldTheory.Finite.Basic
public import Mathlib.FieldTheory.Galois.Basic
public import Mathlib.LinearAlgebra.AnnihilatingPolynomial
public import Mathlib.LinearAlgebra.Matrix.Nondegenerate

/-!
# The normal basis theorem

We prove the normal basis theorem `IsGalois.normalBasis`:
every finite Galois extension has a basis that is an orbit under the Galois group action.

The proof follows [ConradLinearChar] Keith Conrad, *Linear Independence of Characters*.

-/

@[expose] public section

variable (K L : Type*) [Field K] [Field L] [Algebra K L]

open Polynomial FiniteField Module Submodule LinearMap in
-- [ConradLinearChar] Theorem 3.7.
private theorem exists_linearIndependent_algEquiv_apply_of_finite [Finite L] :
    ∃ x : L, LinearIndependent K fun σ : Gal(L/K) ↦ σ x := by
  have := Finite.of_injective _ (algebraMap K L).injective
  have := Fintype.ofFinite K
  /- Since `K[X]` is a PID and `L` is a finitely generated `K[X]`-module (with `X` acting as
    the Frobenius automorphism), there exists `x : L` whose `K[X]`-annihilator is equal to the
    `K[X]`-annihilator of `L`, which is generated by `X ^ [L : K] - 1`. -/
  obtain ⟨x, hx⟩ := exists_ker_toSpanSingleton_eq_annihilator K[X]
    (AEval' (frobeniusAlgHom K L).toLinearMap)
  obtain ⟨x, rfl⟩ := (AEval'.of _).surjective x
  use x
  rw [← span_minpoly_eq_annihilator, minpoly_frobeniusAlgHom, eq_comm] at hx
  rw [← linearIndependent_equiv ((finCongr <| natDegree_X_pow_sub_C (R := K)).trans <|
    .ofBijective _ <| bijective_frobeniusAlgEquivOfAlgebraic_pow K L)]
  /- Therefore, `{Frⁱ | 0 ≤ i < [L : K]}` is linearly independent, which implies that
    `{Frⁱ(x) | 0 ≤ i < [L : K]}` is also linearly independent. -/
  convert (AdjoinRoot.powerBasis (X_pow_sub_C_ne_zero Module.finrank_pos 1)).basis.linearIndependent
    |>.map' ((AEval'.of _).symm.toLinearMap ∘ₗ (liftQ _ _ hx.le).restrictScalars K) <| by
    exact congr($(ker_liftQ_eq_bot' _ _ hx).restrictScalars K)
  ext i
  simp_rw [Equiv.coe_trans, Function.comp_apply, finCongr_apply, Equiv.ofBijective_apply,
    AlgEquiv.coe_pow, AdjoinRoot.powerBasis, AdjoinRoot.powerBasisAux, Basis.coe_mk,
    coe_comp, LinearEquiv.coe_coe, LinearMap.coe_restrictScalars]
  refine (LinearEquiv.eq_symm_apply _).mpr (.symm <| (liftQ_apply ..).trans ?_)
  rw [toSpanSingleton_apply, AEval'.X_pow_smul_of, End.smul_def, End.coe_pow]
  rfl

variable [FiniteDimensional K L]

-- [ConradLinearChar] Theorem 3.6.
private theorem exists_linearIndependent_algEquiv_apply_of_infinite [Infinite K] :
    ∃ x : L, LinearIndependent K fun σ : Gal(L/K) ↦ σ x := by
  classical
  /- Choose a basis `e` of `L` over `K` and form the matrix `M` with entries
    `∑ₖ i⁻¹j(eₖ)Xₖ` where `i,j ∈ Gal(L/K)` and `Xₖ`s are independent variables. -/
  have e := Module.Free.chooseBasis K L
  let M : Matrix Gal(L/K) Gal(L/K) (MvPolynomial _ L) :=
    .of fun i j ↦ ∑ k, i.symm (j (e k)) • .X k
  /- By [ConradLinearChar] Lemma 3.4, there exists `{xₖ} ⊆ L` such that `∑ₖ j(eₖ)xₖ = 1`
    if `j = 1` and `∑ₖ j(eₖ)xₖ = 0` otherwise. -/
  have hq : Submodule.span L (Set.range fun k (j : Gal(L/K)) ↦ j (e k)) = ⊤ := by
    apply span_flip_eq_top_iff_linearIndependent.mpr <|
      ((linearIndependent_algHom_toLinearMap K L L).comp _
        (algEquivEquivAlgHom K L).injective).map' _ (e.constr L).symm.ker
  /- If we evaluate `M` at `Xₖ = xₖ`, the determinant of `M` evaluates to 1. -/
  obtain ⟨c, hc⟩ : ∃ c : _ → L, M.det.eval c = 1 := by
    obtain ⟨g, hg⟩ := (Submodule.mem_span_range_iff_exists_fun _).mp
      (hq ▸ Submodule.mem_top (x := Pi.single 1 1))
    simp_rw [RingHom.map_det]
    refine ⟨g, congr(Matrix.det $(?_)).trans Matrix.det_one⟩
    ext i j
    simpa [M, Pi.single_apply, inv_mul_eq_one, mul_comm, Matrix.one_apply]
      using congr($hg (i⁻¹ * j))
  /- Therefore `det M` is nonzero. -/
  have hM : M.det ≠ 0 := fun h0 ↦ by
    simpa [hc] using congr(($h0).eval c)
  /- Since `K` is infinite, we may evaluate `det M` at some point with coordinates in `K`
    and get a nonzero value. The coordinates give rise to an element of `L` via the basis `e`. -/
  obtain ⟨b, hb⟩ : ∃ b : _ → K, M.det.eval (algebraMap K L ∘ b) ≠ 0 := by
    by_contra! h
    refine hM (MvPolynomial.funext_set _
      (fun _ ↦ Set.infinite_range_of_injective (algebraMap K L).injective) fun x hx ↦ ?_)
    obtain ⟨x, rfl⟩ := Set.range_piMap _ ▸ hx
    simpa using h x
  /- This element of `L` is exactly what we want: we simply need to show the first row of the
    evaluated matrix is `K`-linearly independent. But since the other rows are obtained from the
    first row by applying a `K`-endomorphism, it suffices to show that the columns are linearly
    independent, which follows from the nonzeroness of the evaluated determinant. -/
  refine ⟨∑ k, b k • e k, Fintype.linearIndependent_iff.mpr fun a ha ↦
    funext_iff.mp <| (algebraMap K L).injective.comp_left ?_⟩
  simp_rw [Function.comp_def, map_zero]
  rw [RingHom.map_det, RingHom.mapMatrix_apply] at hb
  refine Matrix.eq_zero_of_mulVec_eq_zero hb (funext fun i ↦ i.injective ?_)
  simp_rw [M, Pi.zero_apply, map_zero, ← ha]
  simp [Algebra.smul_def, Matrix.mulVec_eq_sum, mul_comm]

theorem exists_linearIndependent_algEquiv_apply :
    ∃ x : L, LinearIndependent K fun σ : Gal(L/K) ↦ σ x := by
  obtain h | h := finite_or_infinite K
  · have := Module.finite_of_finite K (M := L)
    exact exists_linearIndependent_algEquiv_apply_of_finite K L
  · exact exists_linearIndependent_algEquiv_apply_of_infinite K L

namespace IsGalois

variable [IsGalois K L]

/-- Given a finite Galois extension `L/K`, `normalBasis K L` is a basis of `L` over `K`
that is an orbit under the Galois group action. -/
noncomputable def normalBasis : Module.Basis Gal(L/K) K L :=
  basisOfLinearIndependentOfCardEqFinrank
    (exists_linearIndependent_algEquiv_apply K L).choose_spec
    (Fintype.card_eq_nat_card.trans <| card_aut_eq_finrank K L)

variable {K L}

theorem normalBasis_apply (e : Gal(L/K)) : normalBasis K L e = e (normalBasis K L 1) := by
  rw [normalBasis, coe_basisOfLinearIndependentOfCardEqFinrank, AlgEquiv.one_apply]

end IsGalois
